from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()


from utils.logger_config import LoggerSingleton
logger = LoggerSingleton.get_logger() 

from langgraph.types import Command



# 定义图
from typing_extensions import TypedDict
from typing import Annotated
from langgraph.graph.message import add_messages
from langgraph.graph import StateGraph, START, END

## 初始化原始图
class State(TypedDict):
    messages: Annotated[list, add_messages]
graph_builder = StateGraph(State)

## 导入工具和模型
from tools_fc.Multiply import multiply_tool
from tools_fc.Add import add_tool
from tools_fc.Human_assistance import human_assistance_tool

tools = [multiply_tool,add_tool,human_assistance_tool]
from model_init.get_model_ollama import ollama_model
llm_with_tools = ollama_model.bind_tools(tools)

## 创建模型和工具节点
### 模型
def chatbot(state: State):
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

graph_builder.add_node("chatbot", chatbot)

### 工具
import json
from langchain_core.messages import ToolMessage
from langgraph.prebuilt import ToolNode, tools_condition


tool_node = ToolNode(tools=tools)
graph_builder.add_node("tools", tool_node)

# 如果聊天机器人要求使用工具，那么‘ tools_condition ’函数返回“tools”，如果可以直接响应，则返回“END”。这个条件路由定义了主代理循环。
graph_builder.add_conditional_edges(
    "chatbot",
    tools_condition,
)
# 任何时候调用一个工具，我们都会返回到聊天机器人来决定下一步
graph_builder.add_edge("tools", "chatbot")
graph_builder.add_edge(START, "chatbot")
graph = graph_builder.compile(checkpointer=memory)

# 绘图——可视化
flag_visual = 0
if flag_visual != 0:
    logger.info('graph可视化')
    from utils.visual_drawing import save_graph_image
    save_graph_image(graph)
else:
    logger.info('graph未可视化')

config = {"configurable": {"thread_id": "1"}}


# user_input = "我需要一些专家的指导来建立一个人工智能代理。你能帮我请求帮助吗？"
# config = {"configurable": {"thread_id": "1"}}

# events = graph.stream(
#     {"messages": [{"role": "user", "content": user_input}]},
#     config,
#     stream_mode="values",
# )
# for event in events:
#     if "messages" in event:
#         event["messages"][-1].pretty_print()

# snapshot = graph.get_state(config)
# snapshot.next


# human_response = (
#     "我们，专家在这里提供帮助！我们建议您使用LangGraph来构建代理。"
#     " 它比简单的自治代理更加可靠和可扩展。"
# )

# human_command = Command(resume={"data": human_response})

# events = graph.stream(human_command, config, stream_mode="values")
# for event in events:
#     if "messages" in event:
#         event["messages"][-1].pretty_print()

def stream_graph_updates(user_input: str):
    events = graph.stream(
        {"messages": [{"role": "user", "content": user_input}]},
        config,
        stream_mode="values",
    )
    for event in events:
        event["messages"][-1].pretty_print()

while True:
    try:
        user_input = input("User(quit\exit\q 退出): ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break

        stream_graph_updates(user_input)
    except:
        # 如果input() 不可用，则回退
        user_input = "What do you know about LangGraph?"
        print("User: " + user_input)
        stream_graph_updates(user_input)
        break